// Copyright (C) ustc9703
// All rights reserved.
// 
// Author : Jiang Li
// Date   : 2004-12-19

#include "stdafx.h"
#include "base_process.h"
#include <winternl.h>
#include <psapi.h>

namespace JoyServer
{


#define SystemProcessesAndThreadsInformation	5

//##ModelId=44CEFAD101AE
typedef struct _CLIENT_ID {
    DWORD	    UniqueProcess;
    DWORD	    UniqueThread;
} CLIENT_ID, * PCLIENT_ID;

//##ModelId=44CEFAD101B8
typedef struct _VM_COUNTERS {
    SIZE_T	    PeakVirtualSize;
    SIZE_T	    VirtualSize;
    ULONG	    PageFaultCount;
    SIZE_T	    PeakWorkingSetSize;
    SIZE_T	    WorkingSetSize;
    SIZE_T	    QuotaPeakPagedPoolUsage;
    SIZE_T	    QuotaPagedPoolUsage;
    SIZE_T	    QuotaPeakNonPagedPoolUsage;
    SIZE_T	    QuotaNonPagedPoolUsage;
    SIZE_T	    PagefileUsage;
    SIZE_T	    PeakPagefileUsage;
} VM_COUNTERS;

//##ModelId=44CEFAD101C6
typedef ULONG KPRIORITY;

//##ModelId=44CEFAD101C8
typedef struct _SYSTEM_THREAD_INFORMATION {
    LARGE_INTEGER   KernelTime;
    LARGE_INTEGER   UserTime;
    LARGE_INTEGER   CreateTime;
    ULONG			WaitTime;
    PVOID			StartAddress;
    CLIENT_ID	    ClientId;
    KPRIORITY	    Priority;
    KPRIORITY	    BasePriority;
    ULONG			ContextSwitchCount;
    LONG			State;
    LONG			WaitReason;
} SYSTEM_THREAD_INFORMATION, * PSYSTEM_THREAD_INFORMATION;

// Note, that the size of the SYSTEM_PROCESS_INFORMATION structure is 
// different on NT 4 and Win2K.

//##ModelId=44CEFAD101D7
typedef struct _SYSTEM_PROCESS_INFORMATION_NT4 {
    ULONG			NextEntryDelta;
    ULONG			ThreadCount;
    ULONG			Reserved1[6];
    LARGE_INTEGER   CreateTime;
    LARGE_INTEGER   UserTime;
    LARGE_INTEGER   KernelTime;
    UNICODE_STRING  ProcessName;
    KPRIORITY	    BasePriority;
    ULONG			ProcessId;
    ULONG			InheritedFromProcessId;
    ULONG			HandleCount;
    ULONG			Reserved2[2];
    VM_COUNTERS	    VmCounters;
    SYSTEM_THREAD_INFORMATION  Threads[1];
} SYSTEM_PROCESS_INFORMATION_NT4, * PSYSTEM_PROCESS_INFORMATION_NT4;

//##ModelId=44CEFAD101F4
typedef struct _SYSTEM_PROCESS_INFORMATION {
    ULONG			NextEntryDelta;
    ULONG			ThreadCount;
    ULONG			Reserved1[6];
    LARGE_INTEGER   CreateTime;
    LARGE_INTEGER   UserTime;
    LARGE_INTEGER   KernelTime;
    UNICODE_STRING  ProcessName;
    KPRIORITY	    BasePriority;
    ULONG			ProcessId;
    ULONG			InheritedFromProcessId;
    ULONG			HandleCount;
    ULONG			Reserved2[2];
    VM_COUNTERS	    VmCounters;
    IO_COUNTERS	    IoCounters;
    SYSTEM_THREAD_INFORMATION  Threads[1];
} SYSTEM_PROCESS_INFORMATION, * PSYSTEM_PROCESS_INFORMATION;

#define OBJ_INHERIT             0x00000002L
#define OBJ_PERMANENT           0x00000010L
#define OBJ_EXCLUSIVE           0x00000020L
#define OBJ_CASE_INSENSITIVE    0x00000040L
#define OBJ_OPENIF              0x00000080L
#define OBJ_OPENLINK            0x00000100L
#define OBJ_KERNEL_HANDLE       0x00000200L
#define OBJ_VALID_ATTRIBUTES    0x000003F2L

//##ModelId=44CEFAD10203
typedef struct _OBJECT_ATTRIBUTES {
    ULONG Length;
    HANDLE RootDirectory;
    PUNICODE_STRING ObjectName;
    ULONG Attributes;
    PVOID SecurityDescriptor;
    PVOID SecurityQualityOfService;
} OBJECT_ATTRIBUTES, * POBJECT_ATTRIBUTES;

NTSYSAPI
NTSTATUS
NTAPI
ZwQuerySystemInformation(
    IN UINT SystemInformationClass,
    IN OUT PVOID SystemInformation,
    IN ULONG SystemInformationLength,
    OUT PULONG ReturnLength OPTIONAL
    );

//##ModelId=44CEFAD10213
typedef NTSTATUS
(WINAPI *PNtQuerySystemInformation) (
    IN UINT SystemInformationClass,
    IN OUT PVOID SystemInformation,
    IN ULONG SystemInformationLength,
    OUT PULONG ReturnLength OPTIONAL
    );

static PNtQuerySystemInformation GetNtQuerySystemInformation()
{
    HMODULE hMod = GetModuleHandleW(L"ntdll.dll");
    if(hMod == NULL)
    {
        return NULL;
    }

    PNtQuerySystemInformation proc = (PNtQuerySystemInformation)
        GetProcAddress(hMod, "NtQuerySystemInformation");

    return proc;
}

static NTSTATUS
MyNtQuerySystemInformation (
    IN UINT SystemInformationClass,
    IN OUT PVOID SystemInformation,
    IN ULONG SystemInformationLength,
    OUT PULONG ReturnLength OPTIONAL
    )
{
    static PNtQuerySystemInformation proc = GetNtQuerySystemInformation();

    if(proc == NULL)
    {
        return -1;
    }

    NTSTATUS sta = proc(
        SystemInformationClass,
        SystemInformation,
        SystemInformationLength,
        ReturnLength);

    return sta;
}

//##ModelId=44CEFAD10215
typedef NTSTATUS
(WINAPI *PNtQueryInformationProcess) (
    IN HANDLE ProcessHandle,
    IN PROCESSINFOCLASS ProcessInformationClass,
    OUT PVOID ProcessInformation,
    IN ULONG ProcessInformationLength,
    OUT PULONG ReturnLength OPTIONAL
    );

static PNtQueryInformationProcess GetNtQueryInformationProcess()
{
    HMODULE hMod = GetModuleHandleW(L"ntdll.dll");
    if(hMod == NULL)
    {
        return NULL;
    }

    PNtQueryInformationProcess proc = (PNtQueryInformationProcess)
        GetProcAddress(hMod, "NtQueryInformationProcess");

    return proc;
}

static NTSTATUS
MyNtQueryInformationProcess (
    IN HANDLE ProcessHandle,
    IN PROCESSINFOCLASS ProcessInformationClass,
    OUT PVOID ProcessInformation,
    IN ULONG ProcessInformationLength,
    OUT PULONG ReturnLength OPTIONAL
    )
{
    static PNtQueryInformationProcess proc = GetNtQueryInformationProcess();

    if(proc == NULL)
    {
        return -1;
    }

    NTSTATUS sta = proc(
        ProcessHandle,
        ProcessInformationClass,
        ProcessInformation,
        ProcessInformationLength,
        ReturnLength);

    return sta;
}

UINT GetOpenedHandleCount()
{
    UINT c = 0;
    MyNtQueryInformationProcess(
        GetCurrentProcess(),
        (PROCESSINFOCLASS)0x14, // unknown value. find it from kernel32.dll.
        &c,
        sizeof(UINT),
        NULL);
    return c;
}

UINT GetThreadCount()
{
    /*
    PERFORMANCE_INFORMATION a;
    GetPerformanceInfo(&a, sizeof(a));
    return a.ThreadCount;
    */
    BYTE buf[128*1024];
    ULONG len = 0;
    DWORD processId = GetCurrentProcessId();
    MyNtQuerySystemInformation(
        SystemProcessInformation,
        buf, sizeof(buf), &len);

    SYSTEM_PROCESS_INFORMATION* infos = (SYSTEM_PROCESS_INFORMATION*)buf;
    SYSTEM_PROCESS_INFORMATION* infoe = (SYSTEM_PROCESS_INFORMATION*)(buf+len);
    for(; infos < infoe; infos = (SYSTEM_PROCESS_INFORMATION*)((char*)infos+infos->NextEntryDelta))
    {
        if(processId == infos->ProcessId)
        {
            return infos->ThreadCount;
        }
    }

    return 0;
}

} // namespace JoyServer
